import os
import subprocess
import re
import multiprocessing

search_bash_name = "search.bash"
download_bash_name = "download.bash"
save_name = "info.txt"

download_dir_prefix = "."

variables = ["huss", "pr", "psl", "tas", "tasmax", "tasmin"]
experiments = ["historical", "ssp126", "ssp245", "ssp370", "ssp585"]
frequency = "day"
skips = ["Amon"]
models = ["AWI-CM-1-1-MR"]
realization = "1"


def download_bash(model, variable, experiment, download_dir):
    args = ["~/.local/bin/acccmip6", 
        "-o", "D" if download_dir != None else "S",
        "-m", model, # model/source
        "-v", variable, # variable
        "-f", frequency, # frequency
        "-e", experiment, # experiment,
        "-rlzn", realization
    ]

    #skip
    if len(skips) != 0: 
        args.append("-skip")
        args.append("\'" + str(skips) + "\'")

    if download_dir != None:
        args.append("-dir")
        args.append(download_dir)

    return args

def create_env():
    for model in models:
        model_dir = download_dir_prefix + "/" + model
        for experiment in experiments:
            experiment_dir = model_dir + "/" + experiment
            os.makedirs(experiment_dir)
            with open(experiment_dir + "/" + download_bash_name, "w") as fe:
                for variable in variables:
                    variable_dir = experiment_dir + "/" + variable
                    os.makedirs(variable_dir)
                    fe.write(" ".join(download_bash(model, variable, experiment, "./" + variable)) + '\n')
                    with open(variable_dir + "/" + download_bash_name, "w") as f:
                        f.write(" ".join(download_bash(model, variable, experiment, ".")))


class MyProcess(multiprocessing.Process):
    def __init__(self, cmd, path, q):
        multiprocessing.Process.__init__(self)
        self.cmd = cmd
        self.path = path
        self.q = q
    def run(self):
        result = subprocess.run(
            self.cmd,
            cwd=self.path,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        output = str(result.stdout, encoding="utf-8")
        error = str(result.stderr, encoding="utf-8")
        self.q.put((self.path, output, error))
        # print(self.output)
        # print(self.error)

def search_data():
    ps = []
    q = multiprocessing.Queue()
    pattern = r'Number of files:[\x1b]\[0m (\d+)'
    for model in models:
        model_dir = download_dir_prefix + "/" + model
        for experiment in experiments:
            experiment_dir = model_dir + "/" + experiment
            for variable in variables:
                variable_dir = experiment_dir + "/" + variable
                p = MyProcess(" ".join(download_bash(model, variable, experiment, None)), variable_dir, q)
                ps.append(p)
                p.start()
    for _ in ps:
        while(q.empty()):
            continue
        path, output, _ = q.get()
        search_result = re.search(pattern, output)
        if search_result != None:
            count = int(search_result.group(1))
        else:
            count = 0
        num = len([f for f in os.listdir(path) if f.endswith(".nc")])
        with open(path + "/" + save_name, "w") as f:
            f.write(path + " " + str(count) + "\n")
            f.write("count:" + str(num))
            if (count == num):
                f.write("\nDownload finished\n")
        print(path, count)

                
if __name__ =="__main__" :
    create_env()
    # search_data()